import argparse
from mcp.server.fastmcp import FastMCP
from typing import Dict, List, Union
from logger import get_logger

# 获取logger
logger = get_logger('server')

# 定义全局变量
# 全局配置
config = {
    'comfyui_ip': 'http://127.0.0.1:8188'  # ComfyUI 默认地址
}

# 初始化MCP
mcp = FastMCP("GBoxComfyUI")
# 启用工具变更通知
mcp._mcp_server.notification_options.tools_changed = True

@mcp.tool()
def generate_ai_image(prompt: str, workflow_file: str = None, output_dir: str = None) -> Dict[str, Union[List[str], List[str]]]:
    """使用 ComfyUI 根据文本提示生成 AI 图像。
    
    使用场景：
    - 当你需要根据文本描述生成图像时
    - 为项目创建视觉内容或概念图
    - 生成艺术作品或设计灵感
    - 根据用户输入的描述创建可视化内容
    
    Args:
        prompt: 描述想要生成的图像内容的ComfyUI 提示词,如果用户没有特殊要求,需要帮助用户优化提示词，不需要背景透明之类的提示词.
        workflow_file: 可选，指定自定义的 ComfyUI 工作流文件路径
        output_dir: 可选，指定输出的完整目录路径，默认使用系统临时目录
        
    Returns:
        包含生成图像信息的字典，包括：
        - file_paths: 保存的图像文件路径列表
        
    注意:
        生成完成后，告诉用户生成图片的位置，同时帮用户在当前编辑器里打开图片
    """
    from comfyui.comfyui import ComfyUI
    
    try:
        # 创建ComfyUI实例，使用配置的地址
        comfy = ComfyUI(config['comfyui_ip'])
        file_paths = comfy.generate_images(
            prompt_text=prompt,
            workflow_path=workflow_file,
            output_dir=output_dir
        )
        # 返回结果
        return {
            "file_paths": file_paths
        }
    except Exception as e:
        logger.error(f"生成图像时发生错误: {str(e)}")
        return {"error": str(e), "file_paths": []}

def parse_args():
    """解析命令行参数并更新全局配置"""
    parser = argparse.ArgumentParser(description='GBox MCP服务器')
    parser.add_argument('--comfyui_ip', type=str, help='ComfyUI 服务器地址')
    
    # 解析参数，只更新config中已有的项
    args = parser.parse_args()
    for k, v in vars(args).items():
        if k in config and v is not None:
            config[k] = v
    
    logger.info(f"使用服务器配置: " + ", ".join(f"{k}={v}" for k, v in config.items()))
    return config

if __name__ == "__main__":
    # 解析命令行参数并更新配置
    params = parse_args()
    mcp.run(transport='stdio')
    
    